package com.jboltai.util.other;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class ExpiringCache<K, V> {
    /**
     * Inner class to store a value with its expiration timestamp
     */
    private static class CacheEntry<V> {
        private final V value;
        private final long expirationTime;

        public CacheEntry(V value, long expirationTime) {
            this.value = value;
            this.expirationTime = expirationTime;
        }

        public boolean isExpired() {
            return System.currentTimeMillis() > expirationTime;
        }

        public V getValue() {
            return value;
        }
    }

    private final Map<K, CacheEntry<V>> cache;
    private final ScheduledExecutorService scheduler;
    private final long defaultTTL; // Time-To-Live in milliseconds

    /**
     * Creates a new ExpiringCache with the specified default time-to-live.
     *
     * @param defaultTTL Default time-to-live in milliseconds
     */
    public ExpiringCache(long defaultTTL) {
        this.defaultTTL = defaultTTL;
        this.cache = new ConcurrentHashMap<>();
        this.scheduler = Executors.newScheduledThreadPool(1);

        // Schedule periodic cleanup task
        scheduler.scheduleAtFixedRate(this::cleanupExpiredEntries,
                defaultTTL / 2, defaultTTL / 2, TimeUnit.MILLISECONDS);
    }

    /**
     * Put a value in the cache with the default TTL.
     *
     * @param key The key
     * @param value The value
     */
    public void put(K key, V value) {
        put(key, value, defaultTTL);
    }

    /**
     * Put a value in the cache with a specific TTL.
     *
     * @param key The key
     * @param value The value
     * @param ttl Time-to-live in milliseconds
     */
    public void put(K key, V value, long ttl) {
        long expirationTime = System.currentTimeMillis() + ttl;
        cache.put(key, new CacheEntry<>(value, expirationTime));
    }

    /**
     * Get a value from the cache.
     *
     * @param key The key
     * @return The value, or null if not found or expired
     */
    public V get(K key) {
        CacheEntry<V> entry = cache.get(key);

        if (entry == null) {
            return null;
        }

        // Check if the entry is expired
        if (entry.isExpired()) {
            cache.remove(key);
            return null;
        }

        return entry.getValue();
    }

    /**
     * Remove a value from the cache.
     *
     * @param key The key to remove
     * @return The removed value, or null if not found
     */
    public V remove(K key) {
        CacheEntry<V> entry = cache.remove(key);
        return entry != null ? entry.getValue() : null;
    }

    /**
     * Check if the cache contains a non-expired key.
     *
     * @param key The key to check
     * @return true if the key exists and has not expired
     */
    public boolean containsKey(K key) {
        CacheEntry<V> entry = cache.get(key);

        if (entry == null) {
            return false;
        }

        if (entry.isExpired()) {
            cache.remove(key);
            return false;
        }

        return true;
    }

    /**
     * Get the current size of the cache (including expired entries that haven't been cleaned up yet).
     *
     * @return The cache size
     */
    public int size() {
        return cache.size();
    }

    /**
     * Clear all entries from the cache.
     */
    public void clear() {
        cache.clear();
    }

    /**
     * Cleanup expired entries from the cache.
     * This method is called automatically by the scheduler.
     */
    private void cleanupExpiredEntries() {
        long currentTime = System.currentTimeMillis();
        cache.entrySet().removeIf(entry -> currentTime > entry.getValue().expirationTime);
    }

    /**
     * Shutdown the cache cleaner executor service.
     * Should be called when the cache is no longer needed.
     */
    public void shutdown() {
        scheduler.shutdown();
    }
}
